import os
import cv2 as cv
import numpy as np
from sklearn.model_selection import train_test_split


def get_pic_data(dir, IMG_H=100, IMG_W=10):
    files = os.listdir(dir)
    x = []
    for file in files:
        path = os.path.join(dir, file)
        img = cv.imread(path, cv.IMREAD_COLOR)
        img = cv.resize(img, (IMG_H, IMG_W))
        img = img.astype(np.float32) / 255.
        x.append(img)
    x = np.float32(x)
    return x


def load_data(data_path, rand_stat=1, IMG_H=256, IMG_W=256, cat_dir='cat', dog_dir='dog', ch_first=True):
    print('Loading data ...')
    x_cat = get_pic_data(os.path.join(data_path, cat_dir), IMG_H, IMG_W)
    print('x_cat:', x_cat.shape)
    n_cat = len(x_cat)
    x_dog = get_pic_data(os.path.join(data_path, dog_dir), IMG_H, IMG_W)
    print('x_dog:', x_dog.shape)
    n_dog = len(x_dog)
    x = np.concatenate([x_cat, x_dog], axis=0)
    if ch_first:
        x = np.transpose(x, [0, 3, 1, 2])
    y_cat = np.full([n_cat], 0, dtype=np.int32)
    print('y_cat:', y_cat.shape)
    y_dog = np.full([n_dog], 1, dtype=np.int32)
    print('y_dog:', y_dog.shape)
    y = np.concatenate([y_cat, y_dog], axis=0)
    print('x:', x.shape)
    print('y:', y.shape)
    x_train, x_val_test, y_train, y_val_test = train_test_split(x, y, train_size=0.8, random_state=rand_stat, shuffle=True)
    x_val, x_test, y_val, y_test = train_test_split(x_val_test, y_val_test, train_size=0.5, random_state=rand_stat, shuffle=True)
    print('x_train', x_train.shape)
    print('x_val', x_val.shape)
    print('x_test', x_test.shape)
    print('y_train', y_train.shape)
    print('y_val', y_val.shape)
    print('y_test', y_test.shape)
    return (x_train, y_train), (x_val, y_val), (x_test, y_test)
